// Dustin Escoffery
// Copyright 2010
#define WIN32_LEAN_AND_MEAN
#include "hook_winsock.h"
#include "hook_window.h"
#include "hook_ddraw.h"
#include "hook_storm.h"
#include "pehlpapi.h"

static HMODULE hModDDRAW = 0;
static HMODULE hModWinsock = 0;

// Install DirectDraw and Winsock Hooks
FARPROC WINAPI MyGetProcAddress(HMODULE hModule, LPCSTR lpProcName)
{
   FARPROC Proc = GetProcAddress(hModule, lpProcName);
   if (hModule == hModDDRAW)
   {
      if (!strcmp(lpProcName, "DirectDrawCreate"))
      {
         Proc = (FARPROC)MyDirectDrawCreate;
      }
   }
   if (hModule == hModWinsock)
   {
      switch ((int)lpProcName)
      {
         case 2: Proc = (FARPROC)MyBind; break;
         case 20: Proc = (FARPROC)MySendTo; break;
      }
   }
   return Proc;
}

// Find instances of DirectDraw and Winsock
HMODULE WINAPI MyLoadLibrary(LPCTSTR lpFileName)
{
   HMODULE hMod = LoadLibrary(lpFileName);
   if (!strcmp(lpFileName, "ddraw.dll")) hModDDRAW = hMod;
   if (!strcmp(lpFileName, "wsock32.dll")) hModWinsock = hMod;
   if (strstr(lpFileName, "standard.snp"))
   {
      WriteImportAddress(hMod, "KERNEL32.dll", "GetProcAddress", (PROC)MyGetProcAddress);
      WriteImportAddress(hMod, "KERNEL32.dll", "LoadLibraryA", (PROC)MyLoadLibrary);
   }
   if (strstr(lpFileName, "battle.snp"))
   {
      WriteImportAddress(hMod, "WSOCK32.dll", "connect", (PROC)MyConnect);
   }
   return hMod;
}

extern "C" BOOL APIENTRY
DllMain(HMODULE hModule, DWORD dwReason, LPVOID lpReserved)
{
   HMODULE mod = GetModuleHandle(NULL);
   HMODULE storm = GetModuleHandle("storm.dll");
   switch (dwReason)
   {
      case DLL_PROCESS_ATTACH:
         // ddraw
         WriteImportAddress(mod, "KERNEL32.dll", "LoadLibraryA", (PROC)MyLoadLibrary);
         WriteImportAddress(mod, "KERNEL32.dll", "GetProcAddress", (PROC)MyGetProcAddress);
         // window
         WriteImportAddress(mod, "USER32.dll", "CreateWindowExA", (PROC)MyCreateWindowEx);
         WriteImportAddress(mod, "USER32.dll", "RegisterClassExA", (PROC)MyRegisterClassEx);
         WriteImportAddress(mod, "KERNEL32.dll", "CreateEventA", (PROC)MyCreateEvent);
         WriteImportAddress(mod, "USER32.dll", "SetCursor", (PROC)MySetCursor);
         WriteImportAddress(mod, "USER32.dll", "SetCursorPos", (PROC)MySetCursorPos);
         WriteImportAddress(mod, "USER32.dll", "GetCursorPos", (PROC)MyGetCursorPos);
         WriteImportAddress(mod, "USER32.dll", "ClipCursor", (PROC)MyClipCursor);
         WriteImportAddress(mod, "USER32.dll", "IsIconic", (PROC)MyIsIconic);
         // storm (bnet)
         WriteImportAddress(storm, "USER32.dll", "CreateWindowExA", (PROC)MyCreateWindowEx2);
         WriteImportAddress(storm, "USER32.dll", "ClientToScreen", (PROC)MyClientToScreen);
         WriteImportAddress(storm, "USER32.dll", "RegisterClassA", (PROC)MyRegisterClass);
         WriteImportAddress(storm, "KERNEL32.dll", "LoadLibraryA", (PROC)MyLoadLibrary);
         WriteImportAddress(storm, "KERNEL32.dll", "GetProcAddress", (PROC)MyGetProcAddress);
         break;
      case DLL_PROCESS_DETACH:
         break;
   }
   return TRUE;
}
